#pragma once

#include "../core/i_alloc.h"
#include "../numeric/i_nullspace.h"
#include "../optimize/i_lm.h"
#include "i_util.h"

namespace idl
{
  /*This routine normalizes the Homography matrix H such that the vector representing H is unitized*/
  template<typename T>
  inline void i_homography_unitize(const T H[9], T Hn[9])
  {
    i_unitize9(H, Hn);
  }
  /*This routine normalizes the Homography matrix H such that the vector representing H is unitized*/
  template<typename T>
  inline void i_homography_unitize(T H[9])
  {
    i_unitize9(H);
  }

  /*This routine runs the denormalization step for 2D homography estimation
  Note that both Tx and Txp are 3x3 upper triangular similarity transformation*/
  template<typename T> inline void i_homography_denormalization(const T Ht[9], const T Tx[9], const T Txp[9], T H[9])
  {
    T Txpi[9], HT[9];
    i_invert_3x3_upper_triangular(Txp, Txpi);
    i_mult_AB_3x3_3x3_w_B_upper_triangular(Ht, Tx, HT);
    i_mult_AB_3x3_3x3_w_A_upper_triangular(Txpi, HT, H);
  }

  /*Constraints the 3x3 homography matrix X such that xp = H*x.
  These functions gives the constraints from the rows of [xp]_x*H*x*/
  template<typename T> inline void i_homography_constraint_r0(const T x[3], const T xp[3], T v[9])
  {
    v[0] = v[1] = v[2] = (T)0.0; 
    v[3] = -(xp[2] * x[0]);
    v[4] = -(xp[2] * x[1]);
    v[5] = -(xp[2] * x[2]);
    v[6] =  (xp[1] * x[0]);
    v[7] =  (xp[1] * x[1]);
    v[8] =  (xp[1] * x[2]);
  }

  template<typename T> inline void i_homography_constraint_r1(const T x[3], const T xp[3], T v[9])
  {
    v[0] =  (xp[2] * x[0]);
    v[1] =  (xp[2] * x[1]);
    v[2] =  (xp[2] * x[2]);
    v[3] = v[4] = v[5] = (T)0.0;
    v[6] = -(xp[0] * x[0]);
    v[7] = -(xp[0] * x[1]);
    v[8] = -(xp[0] * x[2]);
  }

  template<typename T> inline void i_homography_constraint_r2(const T x[3], const T xp[3], T v[9])
  {
    v[0] = -(xp[1] * x[0]);
    v[1] = -(xp[1] * x[1]);
    v[2] = -(xp[1] * x[2]);
    v[3] =  (xp[0] * x[0]);
    v[4] =  (xp[0] * x[1]);
    v[5] =  (xp[0] * x[2]);
    v[6] =  v[7] = v[8] = (T)0.0;
  }

  template<typename T>
  inline void i_homography_constraint(const T x[3], const T xp[3], T v1[9], T v2[9])
  {
    /*constraints from the two rows of [xp]_x*H*x that do not correspond to the largest coordinate of xp*/
    T x0 = i_abs(xp[0]);
    T x1 = i_abs(xp[1]);
    T x2 = i_abs(xp[2]);
    if (x0 >= x1)
    {
      if (x0 >= x2)
      {
        i_homography_constraint_r1(x, xp, v1);
        i_homography_constraint_r2(x, xp, v2);
      }
      else
      {
        i_homography_constraint_r0(x, xp, v1);
        i_homography_constraint_r1(x, xp, v2);
      }
    }
    else
    {
      if (x1 >= x2)
      {
        i_homography_constraint_r0(x, xp, v1);
        i_homography_constraint_r2(x, xp, v2);
      }
      else
      {
        i_homography_constraint_r0(x, xp, v1);
        i_homography_constraint_r1(x, xp, v2);
      }
    }
  }

  /*This routine solves the four point homography matrix:
  Compute homography matrix H such that xp~=Hx given 4 point correspondences x<->xp.
  The points are given as triplets and quadruplets of homogeneous coordinates stored in the rows of x and xp,
  Note that The user is responsible for applying data normalization to the homogeneous coordinates
  to something good.
  If nullspace_method is :
  0, LU is used to find the nullspace, faster but less accurate
  1, SVD is used to find the nullspace, slower but more accurate
  else, always use SVD*/
  template<typename T>
  inline void i_four_point_homography(const T x[12], const T xp[12], T H[9], int null_space_method = 1)
  {
    T A[72]; /*8x9*/
    /*4 point constraints*/
    i_homography_constraint(x,     xp,     A,      A + 9);
    i_homography_constraint(x + 3, xp + 3, A + 18, A + 27);
    i_homography_constraint(x + 6, xp + 6, A + 36, A + 45);
    i_homography_constraint(x + 9, xp + 9, A + 54, A + 63);
    switch (null_space_method)
    {
    case 0:  i_nullspace_lu_8x9_destroy(A, H); break;
    case 1:  i_nullspace_svd_8x9(A, H);        break;
    default: i_nullspace_svd_8x9(A, H); break;
    }
  }

  /*RANSAC solvers:*/

  /*This routine solves the four point homography matrix:
  Compute homography matrix H such that xp~=Hx given 4 point correspondences x<->xp.
  The points are given as triplets and quadruplets of homogeneous coordinates stored in the rows of x and xp,
  Note that the routine is responsible for applying data normalization to the homogeneous coordinates*/
  template<typename T>
  inline void i_four_point_homography_ransac_hypogenfunc(const T x[12], const T xp[12], T H[9])
  {
    T nx[12], nxp[12], Tx[9], Txp[9], Ht[9];
    //normalize input image coordinates to the normalized space;
    i_point_preconditioning4(x,  nx,  Tx);
    i_point_preconditioning4(xp, nxp, Txp);
    //solve the four-point homography problem
    i_four_point_homography(nx, nxp, Ht, 1);
    //denormalization: 
    i_homography_denormalization(Ht, Tx, Txp, H);
  }

  /*Compute the cost - sum of point to point distance in one view (only for inliers) for the Ransac routine
    candidate points are stored in data in homogeneous coordinate system*/
  template<typename T>
  inline void i_homography_ransac_oneview_costfunc(const T H[9], const T* x, const T* xp, int n, int &nr_liner, int *inliers, T &cost, T tol)
  {
    int i, j;
    T hx[3], uxp[3];
    T tolsqr = i_sqr(tol);
    T point_to_point_sqr_distance = (T)0.0;
    nr_liner = 0;
    cost = (T)0.0;
    for (i = 0, j = 0; i < n; i++, j += 3)
    {
      i_mult_Ax_3x3(H, x+j, hx);
      i_homogeneous_unitize3(hx);
      i_copy3(xp + j, uxp);
      i_homogeneous_unitize3(uxp);
      point_to_point_sqr_distance = i_squaresum_diff3(uxp, hx);
      if (point_to_point_sqr_distance < tolsqr)
      {
        inliers[nr_liner++] = i;
        cost += point_to_point_sqr_distance;
      }
    }
  }

  /*Polish solver:*/

  /*This routine optimize the 3x3 homography matrix with a set of 2D<->2D correspondeces using the Levenberg-Marquardt algorithm.
  The input image points {xs, xps} should be in inhomogeneous coordinates.*/
  template<typename T>
  class HomographyOptimizer
  {
  public:
    typedef T data_type;
    static const int _m = 9;
    static const int _d = 2;
    HomographyOptimizer();
    ~HomographyOptimizer();
    void computeCost(const T *c, T *fvec) const;
    void computeJacobian(const T *c, T *J) const;
    int getNumUnknowns() const { return _m; }
    int getNumberResiduals() const { return _d*_n; };
    bool optimize(const T H[9], const T *xs, const T *xps, int n);
    void getHomography(T H[9]) const;
    void setVerbose(bool verbose) { _verbose = verbose; }
    bool isFailed() const { return !_succeeded; }
  protected:
    T _c[_m]; /*9 parameters to model the 3x3 matrix H*/
    int _n;
    const T *_xs;
    const T *_xps;
    bool _succeeded;
    bool _verbose;
  };

  template<typename T>
  HomographyOptimizer<T>::HomographyOptimizer() : _n(0), _xs(NULL), _xps(NULL), _succeeded(false), _verbose(false)
  {
    i_zero(_c, _m);
  }

  template<typename T>
  HomographyOptimizer<T>::~HomographyOptimizer(){}

  template<typename T>
  void HomographyOptimizer<T>::getHomography(T H[9]) const
  {
    i_copy9(_c, H);
  }

  template<typename T>
  bool HomographyOptimizer<T>::optimize(const T H[9]/*initial homography matrix*/, const T* xs, const T* xps, const int n/*# of 2D<->2D correspondeces*/)
  {
    _succeeded = false;
    if (n == 0 || xs == NULL || xps == NULL)
    {
      return false;
    }
    /*alloc error vector fvec*/
    T* fvec = i_alloc<T>(n*_d);
    if (fvec == NULL)
    {
      return false;
    }
    _n = n;
    _xs = xs;
    _xps = xps;
    i_copy9(H, _c);
    LevenbergMarquardt solver;
    solver.setVerbose(_verbose);
    _succeeded = solver.optimize(*this, _c, fvec);
    i_free(fvec);
    return _succeeded;
  }

  /*use single view error for speed consideration*/
  template<typename T>
  void HomographyOptimizer<T>::computeCost(const T *c, T *fvec) const
  {
    const T h0 = c[0];
    const T h1 = c[1];
    const T h2 = c[2];
    const T h3 = c[3];
    const T h4 = c[4];
    const T h5 = c[5];
    const T h6 = c[6];
    const T h7 = c[7];
    const T h8 = c[8];
    T hx0, hx1, hx2;

    const T *cptr_x = _xs;
    const T *cptr_xp = _xps;
    T *ptr_f = fvec;

    for (int i = 0; i < _n; ++i)
    {
      hx0 = h0*cptr_x[0] + h1*cptr_x[1] + h2;
      hx1 = h3*cptr_x[0] + h4*cptr_x[1] + h5;
      hx2 = h6*cptr_x[0] + h7*cptr_x[1] + h8;
      ptr_f[0] = i_div(hx0, hx2) - cptr_xp[0];
      ptr_f[1] = i_div(hx1, hx2) - cptr_xp[1];
      cptr_x += 2;
      cptr_xp += 2;
      ptr_f += _d;
    }
  }

  template<typename T>
  void HomographyOptimizer<T>::computeJacobian(const T *c, T *J) const
  {
    /*J is 2*_n x 9*/
    int i, j_offset = _m * _d;
    const T h0 = c[0];
    const T h1 = c[1];
    const T h2 = c[2];
    const T h3 = c[3];
    const T h4 = c[4];
    const T h5 = c[5];
    const T h6 = c[6];
    const T h7 = c[7];
    const T h8 = c[8];
    T hx0, hx1, hx2, x, y, sf, sfx, sfy;

    const T *cptr_x  = _xs;
    T* ptr_J = J;

    for (i = 0; i < _n; ++i)
    {
      x = cptr_x[0];
      y = cptr_x[1];

      hx0 = h0*x + h1*y + h2;
      hx1 = h3*x + h4*y + h5;
      hx2 = h6*x + h7*y + h8;
      sf = i_rec(hx2*hx2);

      sfx = sf * x;
      sfy = sf * y;
      
      /*dx/dc*/
      ptr_J[0] = hx2 * sfx; //(x * hx2)*sf;
      ptr_J[1] = hx2 * sfy; //(y * hx2)*sf;
      ptr_J[2] = (hx2)*sf;
      ptr_J[3] = (T)0.0;
      ptr_J[4] = (T)0.0;
      ptr_J[5] = (T)0.0;
      ptr_J[6] = -(hx0 * sfx); //(-hx0 * x)*sf;
      ptr_J[7] = -(hx0 * sfy); //(-hx0 * y)*sf;
      ptr_J[8] = -(hx0 * sf);

      /*dy/dc*/
      ptr_J[9] = (T)0.0;
      ptr_J[10] = (T)0.0;
      ptr_J[11] = (T)0.0;
      ptr_J[12] = ptr_J[0];
      ptr_J[13] = ptr_J[1];
      ptr_J[14] = ptr_J[2];
      ptr_J[15] = -(hx1 * sfx);//(-hx1 * x)*sf;
      ptr_J[16] = -(hx1 * sfy);//(-hx1 * y)*sf;
      ptr_J[17] = -(hx1* sf);
      
      cptr_x  += 2; 
      ptr_J += j_offset;
    }
  }

} /* namespace idl */